"""
股票预测模型推理演示脚本
独立运行，专门用于加载已训练模型并进行推理
"""

import torch
import pandas as pd
from torch import nn
import numpy as np
from sklearn.preprocessing import StandardScaler
import os

# 调用计算设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 模型文件名配置
MODEL_PTH_FILE = "model_stock.pth"
DATA_FILE = "processed_stock_data.csv"

# 神经网络模型定义（与训练脚本保持一致）
class NeuralNetwork(nn.Module):
    def __init__(self, input_size):
        super(NeuralNetwork, self).__init__()
        self.network = nn.Sequential(
            nn.Linear(input_size, 128),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 3)  # 输出3个类别：0, 1, 2
        )

    def forward(self, x):
        return self.network(x)

def load_model_and_data():
    """加载模型和数据"""
    # 检查文件是否存在
    if not os.path.exists(MODEL_PTH_FILE):
        print(f"错误：模型文件 {MODEL_PTH_FILE} 不存在！")
        print("请先运行训练脚本生成模型文件")
        return None, None, None
    
    if not os.path.exists(DATA_FILE):
        print(f"错误：数据文件 {DATA_FILE} 不存在！")
        return None, None, None
    
    # 加载数据
    data = pd.read_csv(DATA_FILE)
    features = data.iloc[:, 2:].values.astype(np.float32)
    labels = data['peek'].values.astype(np.int64)
    
    # 标准化特征数据
    scaler = StandardScaler()
    features = scaler.fit_transform(features)
    
    # 获取输入维度并创建模型
    input_size = features.shape[1]
    model = NeuralNetwork(input_size).to(device)
    
    # 加载模型权重
    model.load_state_dict(torch.load(MODEL_PTH_FILE))
    model.eval()
    
    print(f"模型加载成功：{MODEL_PTH_FILE}")
    print(f"数据加载成功：{DATA_FILE}")
    print(f"数据集大小：{len(data)}")
    print(f"特征维度：{input_size}")
    print(f"标签分布：{np.bincount(labels)}")
    
    return model, features, labels

def predict_single_sample(model, features, labels, sample_idx=0):
    """单个样本推理"""
    print(f"\n=== 单个样本推理（样本索引: {sample_idx}） ===")
    
    if sample_idx >= len(features):
        print(f"错误：样本索引 {sample_idx} 超出范围（最大索引: {len(features)-1}）")
        return
    
    # 准备输入数据
    sample_features = torch.tensor(features[sample_idx]).unsqueeze(0).to(device)
    true_label = labels[sample_idx]
    
    # 模型推理
    with torch.no_grad():
        prediction = model(sample_features)
        probabilities = torch.softmax(prediction, dim=1)
        predicted_class = prediction.argmax(1).item()
    
    print(f"真实标签: {true_label}")
    print(f"预测类别: {predicted_class}")
    print(f"预测正确: {'✓' if predicted_class == true_label else '✗'}")
    print(f"各类别概率:")
    for i, prob in enumerate(probabilities[0]):
        print(f"  类别 {i}: {prob.item():.4f} ({prob.item()*100:.2f}%)")
    
    return predicted_class == true_label

def predict_batch(model, features, labels, batch_size=64):
    """批量推理"""
    print(f"\n=== 批量推理（批次大小: {batch_size}） ===")
    
    total_samples = len(features)
    correct_predictions = 0
    class_stats = {0: {'correct': 0, 'total': 0},
                   1: {'correct': 0, 'total': 0},
                   2: {'correct': 0, 'total': 0}}
    
    # 分批处理
    for i in range(0, total_samples, batch_size):
        end_idx = min(i + batch_size, total_samples)
        batch_features = torch.tensor(features[i:end_idx]).to(device)
        batch_labels = labels[i:end_idx]
        
        # 模型推理
        with torch.no_grad():
            predictions = model(batch_features)
            predicted_labels = predictions.argmax(1).cpu().numpy()
        
        # 统计准确率
        for j, (true_label, pred_label) in enumerate(zip(batch_labels, predicted_labels)):
            class_stats[true_label]['total'] += 1
            if true_label == pred_label:
                correct_predictions += 1
                class_stats[true_label]['correct'] += 1
        
        # 显示进度
        if (i // batch_size) % 2 == 0:  # 每2个批次显示一次
            batch_accuracy = np.sum(batch_labels == predicted_labels) / len(batch_labels) * 100
            print(f"批次 {i//batch_size + 1}: 样本 {i+1}-{end_idx}, 准确率: {batch_accuracy:.2f}%")
    
    # 输出总体统计
    overall_accuracy = correct_predictions / total_samples * 100
    print(f"\n=== 批量推理结果统计 ===")
    print(f"总样本数: {total_samples}")
    print(f"总体准确率: {correct_predictions}/{total_samples} = {overall_accuracy:.2f}%")
    
    print(f"\n各类别统计:")
    for class_id, stats in class_stats.items():
        if stats['total'] > 0:
            class_accuracy = stats['correct'] / stats['total'] * 100
            print(f"类别 {class_id}: {stats['correct']}/{stats['total']} = {class_accuracy:.2f}%")
        else:
            print(f"类别 {class_id}: 无数据")
    
    return overall_accuracy

def interactive_prediction(model, features, labels):
    """交互式预测"""
    print(f"\n=== 交互式预测模式 ===")
    print("输入样本索引进行预测，输入 'q' 退出")
    print(f"可用样本索引范围: 0 - {len(features)-1}")
    
    while True:
        user_input = input("\n请输入样本索引 (或 'q' 退出): ").strip()
        
        if user_input.lower() == 'q':
            print("退出交互式预测模式")
            break
        
        try:
            sample_idx = int(user_input)
            predict_single_sample(model, features, labels, sample_idx)
        except ValueError:
            print("错误：请输入有效的数字索引")
        except Exception as e:
            print(f"预测出错：{e}")

def main():
    """主函数"""
    print("=" * 60)
    print("股票预测模型推理演示")
    print("=" * 60)
    
    # 加载模型和数据
    model, features, labels = load_model_and_data()
    if model is None:
        return
    
    # 执行不同类型的推理演示
    while True:
        print(f"\n请选择推理模式:")
        print("1. 单个样本推理")
        print("2. 批量推理")
        print("3. 交互式预测")
        print("4. 退出")
        
        choice = input("请输入选择 (1-4): ").strip()
        
        if choice == '1':
            # 随机选择一个样本进行推理
            import random
            sample_idx = random.randint(0, len(features)-1)
            predict_single_sample(model, features, labels, sample_idx)
            
        elif choice == '2':
            predict_batch(model, features, labels)
            
        elif choice == '3':
            interactive_prediction(model, features, labels)
            
        elif choice == '4':
            print("程序结束")
            break
            
        else:
            print("无效选择，请重新输入")

if __name__ == "__main__":
    main()